import math
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import os
import re
import time
import random
import pandas as pd
import requests
from bs4 import BeautifulSoup
from fake_useragent import UserAgent
from urllib.parse import urljoin


class LiteratureDownloaderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("文献下载工具")
        self.root.geometry("600x400")

        # 设置样式
        self.style = ttk.Style()
        self.style.configure('TFrame', background='#f0f0f0')
        self.style.configure('TLabel', background='#f0f0f0', font=('Arial', 10))
        self.style.configure('TButton', font=('Arial', 10), padding=5)
        self.style.configure('Header.TLabel', font=('Arial', 12, 'bold'))

        # 主框架
        self.main_frame = ttk.Frame(root)
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 标题
        self.header = ttk.Label(self.main_frame, text="PubMed/Sci-Hub文献下载工具", style='Header.TLabel')
        self.header.pack(pady=(0, 20))

        # 输入区域
        self.input_frame = ttk.Frame(self.main_frame)
        self.input_frame.pack(fill=tk.X, pady=5)

        # Excel文件选择
        self.excel_label = ttk.Label(self.input_frame, text="选择Excel文件:")
        self.excel_label.grid(row=0, column=0, sticky=tk.W, padx=5)

        self.excel_path = tk.StringVar()
        self.excel_entry = ttk.Entry(self.input_frame, textvariable=self.excel_path, width=40)
        self.excel_entry.grid(row=0, column=1, padx=5)

        self.browse_button = ttk.Button(self.input_frame, text="浏览", command=self.browse_excel)
        self.browse_button.grid(row=0, column=2, padx=5)

        # DOI输入（可选）
        self.doi_label = ttk.Label(self.input_frame, text="或输入DOI:")
        self.doi_label.grid(row=1, column=0, sticky=tk.W, padx=5, pady=5)

        self.doi_entry = ttk.Entry(self.input_frame, width=40)
        self.doi_entry.grid(row=1, column=1, padx=5, pady=5)

        # 下载选项
        self.option_frame = ttk.Frame(self.main_frame)
        self.option_frame.pack(fill=tk.X, pady=5)

        self.download_option = tk.StringVar(value="pubmed")
        ttk.Radiobutton(
            self.option_frame,
            text="PubMed下载",
            variable=self.download_option,
            value="pubmed"
        ).grid(row=0, column=0, padx=5)

        ttk.Radiobutton(
            self.option_frame,
            text="Sci-Hub下载",
            variable=self.download_option,
            value="scihub"
        ).grid(row=0, column=1, padx=5)

        # 按钮区域
        self.button_frame = ttk.Frame(self.main_frame)
        self.button_frame.pack(fill=tk.X, pady=20)

        self.deal_button = ttk.Button(
            self.button_frame,
            text="处理数据",
            command=self.deal_data,
            style='TButton'
        )
        self.deal_button.pack(side=tk.LEFT, padx=10, ipadx=20)

        self.submit_button = ttk.Button(
            self.button_frame,
            text="提交",
            command=self.submit_task,
            style='TButton'
        )
        self.submit_button.pack(side=tk.LEFT, padx=10, ipadx=20)

        self.download_button = ttk.Button(
            self.button_frame,
            text="下载",
            command=self.start_download,
            style='TButton'
        )
        self.download_button.pack(side=tk.LEFT, padx=10, ipadx=20)

        # 状态显示
        self.status_frame = ttk.Frame(self.main_frame)
        self.status_frame.pack(fill=tk.BOTH, expand=True)

        self.status_text = tk.Text(
            self.status_frame,
            wrap=tk.WORD,
            height=10,
            state=tk.DISABLED,
            font=('Arial', 9)
        )
        self.status_text.pack(fill=tk.BOTH, expand=True)

        self.scrollbar = ttk.Scrollbar(self.status_text)
        self.scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.status_text.config(yscrollcommand=self.scrollbar.set)
        self.scrollbar.config(command=self.status_text.yview)

        # 进度条
        self.progress = ttk.Progressbar(
            self.main_frame,
            orient=tk.HORIZONTAL,
            length=400,
            mode='determinate'
        )
        self.progress.pack(pady=10)

        # 初始化变量
        self.task_queue = []
        self.current_task = None

    def browse_excel(self):
        filepath = filedialog.askopenfilename(
            title="选择Excel文件",
            filetypes=[("Excel文件", "*.xlsx *.xls"), ("所有文件", "*.*")]
        )
        if filepath:
            self.excel_path.set(filepath)
            self.append_status(f"已选择文件: {filepath}")

    def append_status(self, message):
        self.status_text.config(state=tk.NORMAL)
        self.status_text.insert(tk.END, message + "\n")
        self.status_text.see(tk.END)
        self.status_text.config(state=tk.DISABLED)
        self.root.update()

    def deal_data(self):
        input_file = self.excel_path.get()
        df = pd.read_excel(input_file)

        pmids = df['PMID'].tolist()
        titles = df['Title'].tolist()
        pmcids = df['PMCID'].tolist()
        dois = df['DOI'].tolist()

        freeLinkds = []
        downloads = []
        for pmcid in pmcids:
            downloads.append(0)
            if str(pmcid) == 'nan':
                freeLinkds.append('')
            else:
                freeLinkds.append('https://pmc.ncbi.nlm.nih.gov/articles/' + str(pmcid) + '/')


        data = {
            'PMID': pmids,
            'title': titles,
            'DOI': dois,
            'free_link' : freeLinkds,
            'download': downloads
        }



        newDf = pd.DataFrame(data)

        # 将处理后的数据写入新的Excel文件
        output_file = 'output.xlsx'

        newDf.to_excel(output_file, index=False)

        self.append_status(f"数据已成功处理从 {input_file} 读取，并保存到 {output_file}")

    def submit_task(self):
        excel_file = 'output.xlsx'
        doi = self.doi_entry.get()
        option = self.download_option.get()

        if excel_file and os.path.exists(excel_file):
            self.task_queue.append(("excel", excel_file, option))
            self.append_status(f"已添加任务: 从 {excel_file} {option}下载")
        elif doi:
            self.task_queue.append(("doi", doi, option))
            self.append_status(f"已添加任务: DOI {doi} {option}下载")
        else:
            messagebox.showwarning("警告", "请选择Excel文件或输入DOI")



    def start_download(self):
        if not self.task_queue:
            messagebox.showinfo("提示", "没有待处理的任务")
            return

        self.append_status("开始下载任务...")
        self.progress["maximum"] = len(self.task_queue)

        for i, task in enumerate(self.task_queue):
            self.current_task = task
            self.progress["value"] = i + 1
            self.append_status(f"正在处理任务 {i + 1}/{len(self.task_queue)}")

            task_type, content, option = task
            try:
                if task_type == "excel":
                    if option == "pubmed":
                        self.pubmed_download(content)
                    else:
                        self.process_excel_with_scihub(content)
                else:
                    if option == "scihub":
                        result = self.sci_hub(content)
                        self.append_status(f"DOI {content} 下载{'成功' if result else '失败'}")
            except Exception as e:
                self.append_status(f"任务处理失败: {str(e)}")

            self.root.update()

        self.task_queue = []
        self.append_status("所有任务已完成！")
        messagebox.showinfo("完成", "所有下载任务已完成")

    def pubmed_download(self, excel):
        self.append_status('开始PubMed下载.....')
        df = pd.read_excel(excel, engine='openpyxl')

        if 'download' not in df.columns:
            df['download'] = ''

        total = len(df)
        success = 0

        for index, row in df.iterrows():
            try:
                delay = random.uniform(1, 5)
                time.sleep(delay)
                if pd.notna(row['free_link']) and isinstance(row['free_link'], str):
                    title = re.sub(r'[\\/:*?"<>|]', '_', row['title'])
                    p_url = row['free_link']

                    for _ in range(3):
                        try:
                            response = requests.get(
                                p_url,
                                headers={"User-Agent": UserAgent().random},
                                timeout=30
                            )
                            response.raise_for_status()
                            break
                        except Exception as e:
                            self.append_status(f"[{title}] 请求失败，重试中... 错误：{str(e)}")
                            time.sleep(5)
                    else:
                        self.append_status(f"[{title}] 无法获取页面")
                        df.at[index, 'download'] = 0
                        continue

                    soup = BeautifulSoup(response.text, 'html.parser')
                    link_tem = soup.find('a', {'data-ga-label': 'pdf_download_desktop'})
                    if not link_tem:
                        self.append_status(f"[{title}] 未找到PDF链接")
                        df.at[index, 'download'] = 0
                        continue

                    # 获取到pdf链接地址
                    real_link = link_tem.get('href')
                    # real_link = urljoin('https://www.ncbi.nlm.nih.gov', real_link)
                    real_link = urljoin(p_url, real_link)

                    # 判断结尾是否是pdf
                    if not real_link.lower().endswith('.pdf'):
                        self.append_status(f"[{title}] 无效的PDF链接")
                        df.at[index, 'download'] = 0
                        continue

                    headers = {
                        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36 Edg/136.0.0.0",
                        "Accept": "application/pdf, text/html, application/xhtml+xml, application/xml;q=0.9, */*;q=0.8",
                        "Accept-Language": "en-US,en;q=0.5",
                        "Accept-Encoding": "gzip, deflate, br",
                        "Connection": "keep-alive",
                        "Upgrade-Insecure-Requests": "1"
                    }
                    # 创建文件夹
                    os.makedirs('pdf', exist_ok=True)
                    pdf_content, final_url = self.handle_waiting_page(real_link, headers) #requests.get(real_link, headers= headers, stream=True, timeout=60, allow_redirects= True) as r:
                    if pdf_content:
                        with open(f'pdf/{title}.pdf', 'wb') as f:
                            f.write(pdf_content)

                    # 下载完毕 判断大小，如果是2k直接删除
                    file_size = os.path.getsize(f'pdf/{title}.pdf')
                    if file_size > 8192 :
                        self.append_status(f"[{title}] 下载成功！")
                        df.at[index, 'download'] = 1
                        success += 1
                    else:
                        self.append_status('文件大小异常，执行删除')
                        os.remove(f'pdf/{title}.pdf')
                        df.at[index, 'download'] = 0
                else:
                    df.at[index, 'download'] = 0
            except Exception as e:
                self.append_status(f"[{row.get('title', '未知')}] 处理失败: {str(e)}")
                df.at[index, 'download'] = 0

        try:
            output_path = os.path.join(os.path.dirname(excel), 'RX.xlsx')
            df.to_excel(output_path, engine='openpyxl', index=False)
            self.append_status(f'结果已保存到: {output_path}')
        except PermissionError:
            self.append_status('错误：请关闭Excel文件后重试')
        except Exception as e:
            self.append_status(f'保存失败: {str(e)}')

        self.append_status(f'PubMed下载完成！成功 {success}/{total}')

    def handle_waiting_page(response, original_url, headers):
        time.sleep(3)  # 等待 3 秒

        # 重新请求，可能已经可以下载了
        retry_response = requests.get(original_url, headers=headers, allow_redirects=True)
        if retry_response.headers.get('Content-Type', '').lower() == 'application/pdf':
            return retry_response.content, retry_response.url

        return None, None

    def process_excel_with_scihub(self, excel):
        self.append_status('开始通过Sci-Hub处理Excel文件.....')
        df = pd.read_excel(excel, engine='openpyxl')

        if 'download' not in df.columns:
            df['download'] = ''

        total = len(df)
        success = 0

        for index, row in df.iterrows():
            try:
                delay = random.uniform(1, 5)
                time.sleep(delay)

                if pd.notna(row['DOI']) and isinstance(row['DOI'], str):
                    title = re.sub(r'[\\/:*?"<>|]', '_', row['title'])
                    doi = row['DOI']

                    result = self.sci_hub(doi, filename=title)
                    df.at[index, 'download'] = 1 if result else 0
                    success += 1 if result else 0
                else:
                    df.at[index, 'download'] = 0
            except Exception as e:
                self.append_status(f"[{row.get('title', '未知')}] 处理失败: {str(e)}")
                df.at[index, 'download'] = 0

        try:
            output_path = os.path.join(os.path.dirname(excel), 'RX_scihub.xlsx')
            df.to_excel(output_path, engine='openpyxl', index=False)
            self.append_status(f'结果已保存到: {output_path}')
        except PermissionError:
            self.append_status('错误：请关闭Excel文件后重试')
        except Exception as e:
            self.append_status(f'保存失败: {str(e)}')

        self.append_status(f'Sci-Hub处理完成！成功 {success}/{total}')

    def sci_hub(self, doi=None, folder='pdf', filename=None, retry=3):
        os.makedirs(folder, exist_ok=True)
        scihub_domains = [
            "https://www.sci-hub.se/",
            "https://www.sci-hub.st/",
            "https://www.sci-hub.ru/"
        ]
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept-Language': 'en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7'
        }
        for attempt in range(retry):
            try:
                response = requests.get(
                    scihub_domains[attempt % len(scihub_domains)]+doi, headers=headers, stream=True, timeout=30
                )
                soup = BeautifulSoup(response.content, 'html.parser')

                # soup = soup.find('div', {'id': 'minu'})
                # print(soup)
                iframe = soup.find('embed', {'id': 'pdf'})

                if iframe and (src := iframe.get('src', '')):
                    pdf_url = src if src.startswith('http') else f'https:{src}'
                    break

                # download_div = soup.find('div', {'id': 'buttons'})
                # if download_div and (button := download_div.find('button')):
                #     pdf_url = button.get('onclick', '').split("'")[1]
                #     pdf_url = urljoin(scihub_domains[attempt], pdf_url)
                #     break
            except Exception as e:
                self.append_status(f"尝试 {attempt + 1} 失败: {str(e)}")
                time.sleep(5)
        else:
            self.append_status(f"[{doi}] 无法获取下载链接")
            return 0

        try:
            if pdf_url.startswith('//'):
                pdf_url = f'https:{pdf_url}'
            elif pdf_url.startswith('/'):
                pdf_url = urljoin(scihub_domains[0], pdf_url)

            filename = filename or doi.replace('/', '_')
            local_path = os.path.join(folder, f"{filename}.pdf")

            with requests.get(pdf_url, headers=headers, stream=True, timeout=120) as r:
                r.raise_for_status()
                with open(local_path, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            self.append_status(f"[{filename}] 下载成功！")
            return 1
        except Exception as e:
            self.append_status(f"[{filename}] 下载失败: {str(e)}")
            return 0


if __name__ == "__main__":
    root = tk.Tk()
    app = LiteratureDownloaderApp(root)
    root.mainloop()